#include <cmath>
#include <fstream>
#include <iostream>
#include <set>

#include <fbxsdk.h>

#include "LinearSystem.h"
#include "Topology.h"

static bool printDetail = false;

void saveScalarAsVertexColor(KFbxMesh* mesh, const boost::numeric::ublas::vector<double>& scalarVector, double normalizer = 0.5 * M_PI)
{
	KFbxLayerElementVertexColor* vertexColor = NULL;
	for (int l = 0; l < mesh->GetLayerCount(); ++ l)
	{
		vertexColor = mesh->GetLayer(l)->GetVertexColors();
		if (vertexColor)
		{
			break;
		}
	}

	static const double TWO_PI = M_PI * 2.0;

	if (vertexColor)
	{
		KFbxLayerElementArrayTemplate<KFbxColor>& vertexColorArray = vertexColor->GetDirectArray();
		KFbxLayerElementArrayTemplate<int>& vertexColorIndexArray = vertexColor->GetIndexArray();

		const double minimal = *std::min_element(scalarVector.begin(), scalarVector.end());
		const double maxmual = *std::max_element(scalarVector.begin(), scalarVector.end());
		const double range = maxmual - minimal;

		for (int i = 0; i < mesh->GetPolygonCount(); ++ i)
		{
			int a = mesh->GetPolygonVertex(i, 0);
			int b = mesh->GetPolygonVertex(i, 1);
			int c = mesh->GetPolygonVertex(i, 2);

			double scalarA = scalarVector(a);
			double scalarB = scalarVector(b);
			double scalarC = scalarVector(c);

			KFbxColor colorA(0, 0, 0);
			KFbxColor colorB(0, 0, 0);
			KFbxColor colorC(0, 0, 0);
			
			scalarA >= 0 ? colorA.mRed = (scalarA - minimal) / range : colorA.mBlue = (scalarA - minimal) / range;
			scalarB >= 0 ? colorB.mRed = (scalarB - minimal) / range : colorB.mBlue = (scalarB - minimal) / range;
			scalarC >= 0 ? colorC.mRed = (scalarC - minimal) / range : colorC.mBlue = (scalarC - minimal) / range;

			vertexColorArray.SetAt(i * 3 + 0, colorA);
			vertexColorArray.SetAt(i * 3 + 1, colorB);
			vertexColorArray.SetAt(i * 3 + 2, colorC);
		}
	}
}

static double cosineEdge(KFbxMesh* mesh, int vertexIndex, int polygonIndex)
{
	boost::numeric::ublas::vector<double> gaussianCurvatureVector(mesh->GetControlPointsCount());

	int a = mesh->GetPolygonVertex(polygonIndex, 0);
	int b = mesh->GetPolygonVertex(polygonIndex, 1);
	int c = mesh->GetPolygonVertex(polygonIndex, 2);

	KFbxVector4 vertexA = mesh->GetControlPointAt(a);
	KFbxVector4 vertexB = mesh->GetControlPointAt(b);
	KFbxVector4 vertexC = mesh->GetControlPointAt(c);

	double dot = 0;

	if (vertexIndex == a)
	{
		KFbxVector4 aToB = vertexB - vertexA;
		aToB.Normalize();

		KFbxVector4 aToC = vertexC - vertexA;
		aToC.Normalize();

		dot = aToB.DotProduct(aToC);
	}
	else if (vertexIndex == b)
	{
		KFbxVector4 bToA = vertexA - vertexB;
		bToA.Normalize();

		KFbxVector4 bToC = vertexC - vertexB;
		bToC.Normalize();

		dot = bToA.DotProduct(bToC);
	}else if (vertexIndex == c)
	{
		KFbxVector4 cToA = vertexA - vertexC;
		cToA.Normalize();

		KFbxVector4 cToB = vertexB - vertexC;
		cToB.Normalize();

		dot = cToA.DotProduct(cToB);
	}
	else
	{
		assert(0);
		exit(-1);
	}

	return acos(dot);
}

static void calculateGaussianCurvatureFunction(KFbxMesh* mesh, Topology& topo)
{
	boost::numeric::ublas::vector<double> gaussianCurvatureVector(mesh->GetControlPointsCount());

	int vertexCount = mesh->GetControlPointsCount();
	topo.enterRings();
	for (int v = 0; v < vertexCount; ++ v)
	{
		int triangleCountInRing = topo.getPolygonCountInRing();
		for (int t = 0; t < triangleCountInRing; ++ t)
		{
			int polygonIndexInRing = topo.getPolygonIndexInRing(t);

			gaussianCurvatureVector(v) += cosineEdge(mesh, v, polygonIndexInRing);
		}

		gaussianCurvatureVector(v) = 2 * M_PI - gaussianCurvatureVector(v);
	}
	topo.leaveRings();

	saveScalarAsVertexColor(mesh, gaussianCurvatureVector);
}

static void calculatePrincipleDirection(KFbxMesh* mesh, Topology& topo)
{
	boost::numeric::ublas::vector<double>
}

static void calculateHarmonicFunction(KFbxMesh* mesh, Topology& topo)
{
	// Prepare the sparse matrix
	boost::numeric::ublas::compressed_matrix<double> weightMatrix(mesh->GetControlPointsCount(), mesh->GetControlPointsCount());
	boost::numeric::ublas::compressed_vector<double> harmonicVector(mesh->GetControlPointsCount());
	boost::numeric::ublas::compressed_vector<double> constrainVector(mesh->GetControlPointsCount());

	int vertexCount = mesh->GetControlPointsCount();
	topo.enterRings();
	for (int v = 0; v < vertexCount; ++ v)
	{
		// Diagonal set to 1.
		weightMatrix(v, v) = 1;

		// Computer neighbors.
		std::set<int> vertexIndicesInRing;

		int triangleCountInRing = topo.getPolygonCountInRing();
		topo.beginHalfEdges(triangleCountInRing);
		for (int t = 0; t < triangleCountInRing; ++ t)
		{
			int polygonIndexInRing = topo.getPolygonIndexInRing(t);

			int a = mesh->GetPolygonVertex(polygonIndexInRing, 0);
			int b = mesh->GetPolygonVertex(polygonIndexInRing, 1);
			int c = mesh->GetPolygonVertex(polygonIndexInRing, 2);

			// Filter vertices.
			vertexIndicesInRing.insert(a);
			vertexIndicesInRing.insert(b);
			vertexIndicesInRing.insert(c);

			// Construct the half edges of this ring.
			topo.halfEdge(polygonIndexInRing, a, b, c);
		}
		topo.endHalfEdges();

		vertexIndicesInRing.erase(v);

		// Around v, find each shared edge, get two opposite triangles, calcualte \alpha and \beta, discard the boundary.
		for (std::set<int>::iterator itr = vertexIndicesInRing.begin(); itr != vertexIndicesInRing.end(); ++ itr)
		{
			int middleVertexIndex = *itr;

			if (printDetail)
			{
				std::cout << "[ " << middleVertexIndex << ' ' << v << " ]";
			}

			topo.enterConnection(middleVertexIndex, v);
			if (!topo.atBoundary())
			{
				int leftPolygonIndex = topo.getFirstPolygonIndexSharedHalfEdge();
				int rightPolygonIndex = topo.getSecondPolygonIndexSharedHalfEdge();

				// Get opposite 2 vertices on two triangles.
				int leftOppositeVertexIndex = -1;
				int rightOppositeVertexIndex = -1;

				for(int i = 0; i < 3; ++ i)
				{
					int currentLeftVertexIndex = mesh->GetPolygonVertex(leftPolygonIndex, i);
					int currehtRightVertexIndex = mesh->GetPolygonVertex(rightPolygonIndex, i);

					if (currentLeftVertexIndex != middleVertexIndex && currentLeftVertexIndex != v)
					{
						leftOppositeVertexIndex = currentLeftVertexIndex;
					}

					if (currehtRightVertexIndex != middleVertexIndex && currehtRightVertexIndex != v)
					{
						rightOppositeVertexIndex = currehtRightVertexIndex;
					}
				}
				assert(leftOppositeVertexIndex != -1 && rightOppositeVertexIndex != -1);

				if (printDetail)
				{
					std::cout << "-> [" << leftOppositeVertexIndex << ' ' << rightOppositeVertexIndex <<  " ]" << std::endl;
				}

				// Get cotangent weights to (v, leftOppositeVertexIndex), (v, rightOppositeVertexIndex)
				KFbxVector4 leftU = mesh->GetControlPointAt(v) - mesh->GetControlPointAt(leftOppositeVertexIndex);
				KFbxVector4 leftV = mesh->GetControlPointAt(middleVertexIndex) - mesh->GetControlPointAt(leftOppositeVertexIndex);

				KFbxVector4 rightU = mesh->GetControlPointAt(v) - mesh->GetControlPointAt(rightOppositeVertexIndex);
				KFbxVector4 rightV = mesh->GetControlPointAt(middleVertexIndex) - mesh->GetControlPointAt(rightOppositeVertexIndex);

				double dotLeftUV = leftU.DotProduct(leftV);
				double leftCot = dotLeftUV / sqrt(pow(leftU.Length() * leftV.Length(), 2) - pow(dotLeftUV, 2));
				double dotRightUV = rightU.DotProduct(rightV);
				double rightCot = dotRightUV / sqrt(pow(rightU.Length() * rightV.Length(), 2) - pow(dotRightUV, 2));

				weightMatrix(v, middleVertexIndex) = 0.5 * (leftCot + rightCot);
			}
			else
			{
				if (printDetail)
				{
					std::cout << "-> boundary\n" << std::endl;
				}
			}
			topo.leaveConnection();
		}

		topo.removeHalfEdges();
	}

	topo.leaveRings();

	topo.removeTriangles();

	// Input constrains.
	std::cout << "Input constrain, end with -1.\n";
	int constrainVertexIndex = -1;
	float constrainValue = 0;
	do
	{
		if (constrainVertexIndex > -1 && constrainVertexIndex < mesh->GetControlPointsCount())
		{
			constrainVector(constrainVertexIndex) = constrainValue;
			for (int i = 0; i < mesh->GetControlPointsCount(); ++ i)
			{
				weightMatrix(constrainVertexIndex, i) = 0;
			}
			weightMatrix(constrainVertexIndex, constrainVertexIndex) = 1;
		}
		std::cin >> constrainVertexIndex >> constrainValue;
	}while(constrainVertexIndex != -1);

	// Solve scalar function.
	std::cout << "Solve linear system." << std::endl;
	solveLinearSystem(harmonicVector, weightMatrix, constrainVector);

	saveScalarAsVertexColor(mesh, harmonicVector);
}

void solve(KFbxMesh* mesh)
{
	Topology topo;

	int triangleCount = mesh->GetPolygonCount();
	topo.beginTriangles(triangleCount);
	for (int t = 0; t < triangleCount; ++ t)
	{
		topo.triangle(t, mesh->GetPolygonVertex(t, 0), mesh->GetPolygonVertex(t, 1), mesh->GetPolygonVertex(t, 2));
	}
	topo.endTriangles();

	calculateGaussianCurvatureFunction(mesh, topo);
}

int main(int argc, char* argv[])
{
	-- argc, ++ argv ;
	if (argc != 2)
	{
		return EXIT_FAILURE;
	}

	// Load the scene.
    KFbxSdkManager* sdkManager = KFbxSdkManager::Create();
    KFbxScene* scene = KFbxScene::Create(sdkManager, "");

	KFbxImporter* importer = KFbxImporter::Create(sdkManager, "");
	if (importer->Initialize(argv[0], -1, sdkManager->GetIOSettings()))
	{
		if (!importer->Import(scene))
		{
			return EXIT_FAILURE;
		}
		importer->Destroy();
		importer = NULL;
	}

	KFbxNode* root = scene->GetRootNode();
	KFbxMesh* mesh = NULL;
	for (int i = 0; i < root->GetChildCount(); ++ i)
	{
		KFbxNode* child = root->GetChild(i);
		if (mesh = child->GetMesh())
		{
			break;
		}
	}

	// Solve scalar function on mesh.
	solve(mesh);

	// Write out the result.
	KFbxExporter* exporter = KFbxExporter::Create(sdkManager, "");
	if (exporter->Initialize(argv[1]))
	{
		exporter->Export(scene);
	}
	exporter->Destroy();
	exporter = NULL;

    if (sdkManager)
	{
		sdkManager->Destroy();
		sdkManager = NULL;
	}

	return EXIT_SUCCESS;
}
